#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
系统监控API模块
提供系统资源监控的REST API接口
"""

import os
import sys
import psutil
import torch
import json
from datetime import datetime
from typing import Dict, Any, Optional
from aiohttp import web
import server

# 使用独立的日志系统，不受其他脚本影响
from ..sd_train.lora_trainer_utils.custom_logger import get_logger

# 为这个模块创建独立的logger
logger = get_logger("system_monitor_api")

def get_memory_info():
    memory = psutil.virtual_memory()
    return {
        "total": memory.total,
        "available": memory.available,
        "used": memory.used,
        "percent": memory.percent
    }

def get_gpu_info():
    gpu_info = {}
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        for i in range(torch.cuda.device_count()):
            try:
                props = torch.cuda.get_device_properties(i)
                total_memory = props.total_memory or 0
                reserved_memory = torch.cuda.memory_reserved(i) or 0
                used_memory = torch.cuda.memory_allocated(i) or 0 # 使用allocated更准确反映实际使用

                gpu_info[f"gpu_{i}"] = {
                    "name": props.name,
                    "total_memory": total_memory,
                    "used_memory": used_memory,
                    "reserved_memory": reserved_memory,
                    "percent": (used_memory / total_memory) * 100 if total_memory > 0 else 0,
                }
            except Exception as e:
                logger.warning(f"无法获取GPU {i} 的信息: {e}")
                gpu_info[f"gpu_{i}"] = {
                    "name": f"GPU {i} (Error)",
                    "total_memory": 0, "used_memory": 0, "reserved_memory": 0, "percent": 0
                }
    else:
        # 即使没有GPU，也提供一个默认结构
        gpu_info["gpu_0"] = {
            "name": "N/A",
            "total_memory": 0, "used_memory": 0, "reserved_memory": 0, "percent": 0
        }
    return gpu_info

def get_cpu_info():
    return {
        "percent": psutil.cpu_percent(interval=None), # non-blocking call
        "count": psutil.cpu_count()
    }

@server.PromptServer.instance.routes.get("/system/status")
async def get_system_status(request):
    """获取完整系统状态"""
    try:
        status = {
            "timestamp": datetime.now().isoformat(),
            "memory": get_memory_info(),
            "gpu": get_gpu_info(),
            "cpu": get_cpu_info(),
        }
        return web.json_response(status)
    except Exception as e:
        logger.error(f"获取系统状态失败: {str(e)}", exc_info=True)
        return web.json_response({"error": str(e)}, status=500)

def initialize_system_monitor():
    """
    此函数由 __init__.py 调用。
    API路由已经通过 @server.PromptServer.instance.routes.get 装饰器自动注册。
    此函数保留用于未来可能的其他初始化任务。
    """
    logger.info("✅ 系统监控API路由已通过装饰器方式加载。")

# 在模块加载时调用一次cpu_percent，以避免首次调用返回0
psutil.cpu_percent(interval=None)

def get_routes():
    """获取API路由列表"""
    logger.debug("获取系统监控API路由")
    return server.PromptServer.instance.routes

# __init__.py 会调用这个函数
# 主要的API注册逻辑已经通过 @instance.app.get 装饰器在上面完成了
# 保留这个结构是为了清晰和未来的可扩展性。 